/**
 *    Copyright (C) 2025-present MongoDB, Inc.
 *
 *    This program is free software: you can redistribute it and/or modify
 *    it under the terms of the Server Side Public License, version 1,
 *    as published by MongoDB, Inc.
 *
 *    This program is distributed in the hope that it will be useful,
 *    but WITHOUT ANY WARRANTY; without even the implied warranty of
 *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *    Server Side Public License for more details.
 *
 *    You should have received a copy of the Server Side Public License
 *    along with this program. If not, see
 *    <http://www.mongodb.com/licensing/server-side-public-license>.
 *
 *    As a special exception, the copyright holders give permission to link the
 *    code of portions of this program with the OpenSSL library under certain
 *    conditions as described in each individual source file and distribute
 *    linked combinations including the program with the OpenSSL library. You
 *    must comply with the Server Side Public License in all respects for
 *    all of the code used other than as permitted herein. If you modify file(s)
 *    with this exception, you may extend this exception to your version of the
 *    file(s), but you are not obligated to do so. If you do not wish to do so,
 *    delete this exception statement from your version. If you delete this
 *    exception statement from all source files in the program, then also delete
 *    it in the license file.
 */

#include "mongo/db/query/compiler/optimizer/join/single_table_access.h"

#include "mongo/bson/json.h"
#include "mongo/db/namespace_string.h"
#include "mongo/db/query/compiler/ce/sampling/sampling_test_utils.h"
#include "mongo/db/query/compiler/optimizer/join/unit_test_helpers.h"
#include "mongo/db/query/compiler/physical_model/query_solution/stage_types.h"
#include "mongo/unittest/unittest.h"

namespace mongo::join_ordering {

using namespace mongo::cost_based_ranker;

using SingleTableAccessTestFixture = JoinOrderingTestFixture;

void assertQuerySolutionHasEstimate(const QuerySolutionNode* qsn, const EstimateMap& estimates) {
    auto it = estimates.find(qsn);
    ASSERT(it != estimates.end());
    ASSERT_EQ(EstimationSource::Sampling, it->second.outCE.source());
    for (auto&& child : qsn->children) {
        assertQuerySolutionHasEstimate(child.get(), estimates);
    }
}

// Test estimate map is populated for each collection
TEST_F(SingleTableAccessTestFixture, EstimatesPopulated) {
    auto opCtx = operationContext();
    auto nss1 = NamespaceString::createNamespaceString_forTest("test", "coll1");
    auto nss2 = NamespaceString::createNamespaceString_forTest("test", "coll2");

    std::vector<BSONObj> docs;
    for (int i = 0; i < 10; ++i) {
        docs.push_back(BSON("_id" << i << "a" << 1 << "b" << i));
    }
    ce::createCollAndInsertDocuments(opCtx, nss1, docs);

    // coll2 is 10x larger than coll1.
    for (int i = 10; i < 100; ++i) {
        docs.push_back(BSON("_id" << i << "a" << 1 << "b" << i));
    }
    ce::createCollAndInsertDocuments(opCtx, nss2, docs);

    {
        auto mca = multipleCollectionAccessor(opCtx, {nss1, nss2});
        auto nss1UUID = mca.lookupCollection(nss1)->uuid();

        createIndex(nss1UUID, fromjson("{a: 1}"), "a_1");
        createIndex(nss1UUID, fromjson("{b: 1}"), "b_1");
    }

    // Get new MultiCollectionAccessor after all DDLs are done.
    auto mca = multipleCollectionAccessor(opCtx, {nss1, nss2});

    SamplingEstimatorMap estimators;
    estimators[nss1] = samplingEstimator(mca, nss1);
    estimators[nss2] = samplingEstimator(mca, nss2);

    auto filter1 = fromjson("{a: 1, b: 1}");
    auto filter2 = fromjson("{a: 1}");

    // Mock a JoinGraph for testing purposes.
    JoinGraph graph;
    graph.addNode(nss1, makeCanonicalQuery(nss1, filter1), boost::none);
    auto node2 = graph.addNode(nss2, makeCanonicalQuery(nss2, filter2), boost::none);
    ASSERT(node2);
    auto swRes = singleTableAccessPlans(opCtx, mca, graph, estimators);
    ASSERT_OK(swRes);

    auto& res = swRes.getValue();
    ASSERT_EQ(2, res.solns.size());

    // There are no indexes on nss2, so the chosen access path must use a collection scan.
    auto soln2 = res.solns.at(graph.accessPathAt(*node2)).get();
    ASSERT(soln2);
    ASSERT_EQ(soln2->getFirstNodeByType(STAGE_COLLSCAN).second, 1);

    for (auto&& [_, soln] : res.solns) {
        assertQuerySolutionHasEstimate(soln->root(), res.estimate);
    }
}

}  // namespace mongo::join_ordering
